{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "85a348a5-f139-4cb0-8f82-6ebeabd2c882",
   "metadata": {},
   "source": [
    "# API of implemented methods\n",
    "\n",
    "This notebook spells out the API for all algorithms implemented in the `sbi` toolbox:\n",
    "\n",
    "- Posterior estimation (SNPE)\n",
    "\n",
    "- Likelihood estimation (SNLE)\n",
    "\n",
    "- Likelihood-ratio estimation (SNRE)\n",
    "\n",
    "- Utilities  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e2d1852-898e-49ed-85ef-888f0abd7e4e",
   "metadata": {},
   "source": [
    "## Posterior estimation (SNPE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e7559f4-d97b-4bca-8e22-6529c9c92d9c",
   "metadata": {},
   "source": [
    "**Fast ε-free Inference of Simulation Models with Bayesian Conditional Density Estimation**<br> by Papamakarios & Murray (NeurIPS 2016) <br>[[PDF]](https://papers.nips.cc/paper/6084-fast-free-inference-of-simulation-models-with-bayesian-conditional-density-estimation.pdf) [[BibTeX]](https://papers.nips.cc/paper/6084-fast-free-inference-of-simulation-models-with-bayesian-conditional-density-estimation/bibtex)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "080081fb-1558-4756-8a8f-ff065dd1b400",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import SNPE_A\n",
    "\n",
    "inference = SNPE_A(prior)\n",
    "proposal = prior\n",
    "for _ in range(rounds):\n",
    "    theta = proposal.sample((num_sims,))\n",
    "    x = simulator(theta)\n",
    "    _ = inference.append_simulations(theta, x, proposal=proposal).train()\n",
    "    posterior = inference.build_posterior().set_default_x(x_o)\n",
    "    proposal = posterior"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f04c1c6-35b3-4a91-90c5-9fd0377eb8e2",
   "metadata": {},
   "source": [
    "**Automatic posterior transformation for likelihood-free inference**<br>by Greenberg, Nonnenmacher & Macke (ICML 2019) <br>[[PDF]](http://proceedings.mlr.press/v97/greenberg19a/greenberg19a.pdf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "743eb04f-2cd5-4986-a33b-f2207b9cd5a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import SNPE\n",
    "\n",
    "inference = SNPE(prior)\n",
    "proposal = prior\n",
    "for _ in range(rounds):\n",
    "    theta = proposal.sample((num_sims,))\n",
    "    x = simulator(theta)\n",
    "    _ = inference.append_simulations(theta, x, proposal=proposal).train()\n",
    "    posterior = inference.build_posterior().set_default_x(x_o)\n",
    "    proposal = posterior"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4de2b24-94ce-4cbf-a675-b0c19b5200ca",
   "metadata": {},
   "source": [
    "**Truncated proposals for scalable and hassle-free simulation-based inference** <br> by Deistler, Goncalves & Macke (NeurIPS 2022) <br>[[Paper]](https://arxiv.org/abs/2210.04815)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae54b1a9-c3a6-4ee9-b687-bf8c046023c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import SNPE\n",
    "from sbi.utils import get_density_thresholder, RestrictedPrior\n",
    "\n",
    "inference = SNPE(prior)\n",
    "proposal = prior\n",
    "for _ in range(rounds):\n",
    "    theta = proposal.sample((num_sims,))\n",
    "    x = simulator(theta)\n",
    "    _ = inference.append_simulations(theta, x).train(force_first_round_loss=True)\n",
    "    posterior = inference.build_posterior().set_default_x(x_o)\n",
    "\n",
    "    accept_reject_fn = get_density_thresholder(posterior, quantile=1e-4)\n",
    "    proposal = RestrictedPrior(prior, accept_reject_fn, sample_with=\"rejection\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d13f84e2-d35a-4f54-8cbf-0e4be1a38fb3",
   "metadata": {},
   "source": [
    "## Likelihood estimation (SNLE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89a5783d-1b46-47f8-800b-3cda038ea447",
   "metadata": {},
   "source": [
    "**Sequential neural likelihood: Fast likelihood-free inference with autoregressive flows**<br>by Papamakarios, Sterratt & Murray (AISTATS 2019) <br>[[PDF]](http://proceedings.mlr.press/v89/papamakarios19a/papamakarios19a.pdf) [[BibTeX]](https://gpapamak.github.io/bibtex/snl.bib)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4430dbe-ac60-4978-9695-d0a5b317ee57",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import SNLE\n",
    "\n",
    "inference = SNLE(prior)\n",
    "proposal = prior\n",
    "for _ in range(rounds):\n",
    "    theta = proposal.sample((num_sims,))\n",
    "    x = simulator(theta)\n",
    "    _ = inference.append_simulations(theta, x).train()\n",
    "    posterior = inference.build_posterior().set_default_x(x_o)\n",
    "    proposal = posterior"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16d22866-e095-4f4e-bcef-40bc196a8703",
   "metadata": {},
   "source": [
    "**Variational methods for simulation-based inference** <br> by Glöckler, Deistler, Macke (ICLR 2022) <br>[[Paper]](https://arxiv.org/abs/2203.04176)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d284d6c5-e6f6-4b1d-9c15-d6fa1736a10e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import SNLE\n",
    "\n",
    "inference = SNLE(prior)\n",
    "proposal = prior\n",
    "for _ in range(rounds):\n",
    "    theta = proposal.sample((num_sims,))\n",
    "    x = simulator(theta)\n",
    "    _ = inference.append_simulations(theta, x).train()\n",
    "    posterior = inference.build_posterior(sample_with=\"vi\", vi_method=\"fKL\").set_default_x(x_o)\n",
    "    proposal = posterior"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1ca3f25-c7d2-4688-b5c0-f50864d650ba",
   "metadata": {},
   "source": [
    "**Flexible and efficient simulation-based inference for models of decision-making** <br> by Boelts, Lueckmann, Gao, Macke (Elife 2022) <br>[[Paper]](https://elifesciences.org/articles/77220)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e6d6d8f-8718-44cd-bdf7-a4af2887fc14",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import MNLE\n",
    "\n",
    "inference = MNLE(prior)\n",
    "theta = prior.sample((num_sims,))\n",
    "x = simulator(theta)\n",
    "_ = inference.append_simulations(theta, x).train()\n",
    "posterior = inference.build_posterior().set_default_x(x_o)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d81ff49e-f363-43c0-ba16-3ee6a697be5e",
   "metadata": {},
   "source": [
    "## Likelihood-ratio estimation (SNRE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0823041a-b3f7-4cd3-8a48-497450f622ea",
   "metadata": {},
   "source": [
    "**Likelihood-free MCMC with Amortized Approximate Likelihood Ratios**<br>by Hermans, Begy & Louppe (ICML 2020) <br>[[PDF]](http://proceedings.mlr.press/v119/hermans20a/hermans20a.pdf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b58c3609-7bd7-40ce-a154-f72a190da2ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import SNRE_A\n",
    "\n",
    "inference = SNRE_A(prior)\n",
    "theta = prior.sample((num_sims,))\n",
    "x = simulator(theta)\n",
    "_ = inference.append_simulations(theta, x).train()\n",
    "posterior = inference.build_posterior().set_default_x(x_o)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42b12584-63cd-47d4-93d2-430d407e1e0b",
   "metadata": {},
   "source": [
    "**On Contrastive Learning for Likelihood-free Inference**<br>Durkan, Murray & Papamakarios (ICML 2020) <br>[[PDF]](http://proceedings.mlr.press/v119/durkan20a/durkan20a.pdf)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e36ab4e7-713f-4ff2-b467-8b481a149861",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import SNRE\n",
    "\n",
    "inference = SNRE(prior)\n",
    "proposal = prior\n",
    "for _ in range(rounds):\n",
    "    theta = proposal.sample((num_sims,))\n",
    "    x = simulator(theta)\n",
    "    _ = inference.append_simulations(theta, x).train()\n",
    "    posterior = inference.build_posterior().set_default_x(x_o)\n",
    "    proposal = posterior"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44d0151a",
   "metadata": {},
   "source": [
    "**Towards Reliable Simulation-Based Inference with Balanced Neural Ratio Estimation**<br>by Delaunoy, Hermans, Rozet, Wehenkel & Louppe (NeurIPS 2022) <br>[[PDF]](https://arxiv.org/pdf/2208.13624.pdf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85e6cf8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import BNRE\n",
    "\n",
    "inference = BNRE(prior)\n",
    "theta = prior.sample((num_sims,))\n",
    "x = simulator(theta)\n",
    "_ = inference.append_simulations(theta, x).train(regularization_strength=100.)\n",
    "posterior = inference.build_posterior().set_default_x(x_o)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7ecf9a0-6a82-42e1-9c12-97d0f5d825be",
   "metadata": {},
   "source": [
    "**Contrastive Neural Ratio Estimation**<br>Benjamin Kurt Miller, Christoph Weniger, Patrick Forré (NeurIPS 2022) <br>[[PDF]](https://arxiv.org/pdf/2210.06170.pdf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ec55e76-dd86-46d1-a7cc-643324488820",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The main feature of NRE-C is producing an exact ratio of densities at optimum, even when using multiple contrastive pairs (classes).\n",
    "\n",
    "from sbi.inference import SNRE_C\n",
    "\n",
    "# Amortized inference\n",
    "inference = SNRE_C(prior)\n",
    "proposal = prior\n",
    "theta = proposal.sample((num_sims,))\n",
    "x = simulator(theta)\n",
    "_ = inference.append_simulations(theta, x).train(\n",
    "    num_classes=5,  # SNRE_C sees `2 * num_classes - 1` marginally drawn contrastive pairs.\n",
    "    gamma=1.0,  # SNRE_C can control the weight between terms in its loss function.\n",
    ")\n",
    "posterior = inference.build_posterior().set_default_x(x_o)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6271d3b2-1d64-45b8-93b7-b640ab7dafc5",
   "metadata": {},
   "source": [
    "## Utilities"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07414c0c-f8b5-45bd-9597-0eaf19d50a13",
   "metadata": {},
   "source": [
    "**Simulation-based calibration**<br>by Talts, Betancourt, Simpson, Vehtari, Gelman (arxiv 2018) <br>[[Paper]](https://arxiv.org/abs/1804.06788))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7066ef9b-0e3d-44d3-a80e-5e06de7845ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.analysis import run_sbc, sbc_rank_plot\n",
    "\n",
    "thetas = prior.sample((1_000,))\n",
    "xs = simulator(thetas)\n",
    "\n",
    "ranks, dap_samples = run_sbc(\n",
    "    thetas, xs, posterior, num_posterior_samples=1_000\n",
    ")\n",
    "\n",
    "_ = sbc_rank_plot(\n",
    "    ranks=ranks,\n",
    "    num_posterior_samples=num_posterior_samples,\n",
    "    plot_type=\"hist\",\n",
    "    num_bins=None,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f1c75d9-4139-45ba-ad60-49d8ce8ec2ec",
   "metadata": {},
   "source": [
    "**Restriction estimator**<br>by Deistler, Macke & Goncalves (PNAS 2022) <br>[[Paper]](https://www.pnas.org/doi/10.1073/pnas.2207632119)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdb39216-d0f0-41f6-a065-579d309ee1eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import SNPE\n",
    "from sbi.utils import RestrictionEstimator\n",
    "\n",
    "restriction_estimator = RestrictionEstimator(prior=prior)\n",
    "proposal = prior\n",
    "\n",
    "for _ in range(num_rounds):\n",
    "    theta = proposal.sample((num_sims,))\n",
    "    x = simulator(theta)\n",
    "    restriction_estimator.append_simulations(theta, x)\n",
    "    classifier = restriction_estimator.train()\n",
    "    proposal = restriction_estimator.restrict_prior()\n",
    "\n",
    "all_theta, all_x, _ = restriction_estimator.get_simulations()\n",
    "\n",
    "inference = SNPE(prior)\n",
    "density_estimator = inference.append_simulations(all_theta, all_x).train()\n",
    "posterior = inference.build_posterior()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48853668-7b6f-4cfd-9d93-7d62f0e77de8",
   "metadata": {},
   "source": [
    "**Expected coverage (sample-based)**<br>as computed in Deistler, Goncalves, Macke (Neurips 2022) [[Paper]](https://arxiv.org/abs/2210.04815) and in Rozet, Louppe (2021) [[Paper]](https://matheo.uliege.be/handle/2268.2/12993)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60e3d581-8a7f-4133-8756-9750f0174c88",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.analysis import run_sbc, sbc_rank_plot\n",
    "\n",
    "thetas = prior.sample((1_000,))\n",
    "xs = simulator(thetas)\n",
    "\n",
    "ranks, dap_samples = run_sbc(\n",
    "    thetas, xs, posterior, num_posterior_samples=1_000, reduce_fns=posterior.log_prob\n",
    ")\n",
    "\n",
    "_ = sbc_rank_plot(\n",
    "    ranks=ranks,\n",
    "    num_posterior_samples=num_posterior_samples,\n",
    "    plot_type=\"hist\",\n",
    "    num_bins=None,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "c50aa3a452b5e33eec699c3d0adceaddf116b15627c63bb6b43782d4547b8f5a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
